import os
import rasterio
import numpy as np
from typing import Tuple

def get_LS(sp: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """
    根据坡度计算L和S值。

    参数:
    sp (np.ndarray): 坡度数据数组。

    返回:
    Tuple[np.ndarray, np.ndarray]: 包含L值和S值的元组。
    """
    m = np.where(sp <= 1.0, 2,
            np.where((sp > 1.0) & (sp <= 3.0), 0.3,
            np.where((sp > 3.0) & (sp <= 5.0), 0.4, 0.5)))

    length = 30.0 / np.cos(sp / 180 * np.pi)
    lv = (length / 22.13) ** m

    sv = np.where(sp < 5.0, 10.8 * np.sin(sp / 180 * np.pi) + 0.03,
            np.where((sp >= 5.0) & (sp <= 10.0), 16.8 * np.sin(sp / 180 * np.pi) - 0.5,
            21.9 * np.sin(sp / 180 * np.pi) - 0.96))

    return lv, sv


def calculate_lv_sv(input_file_path: str, output_lv_path: str, output_sv_path: str) -> None:
    """
    从输入的坡度图像文件计算L和S值，并将结果保存到指定的输出文件中。

    参数:
    input_file_path (str): 输入的坡度图像文件路径。
    output_lv_path (str): 输出的L值图像文件路径。
    output_sv_path (str): 输出的S值图像文件路径。
    """
    try:
        with rasterio.open(input_file_path) as src:
            # 读取图像数据
            slop_data = src.read(1)

            # 初始化输出数组
            lv_data = np.zeros(slop_data.shape, dtype=np.float32)
            print("初始化数组lv_data成功！")
            sv_data = np.zeros(slop_data.shape, dtype=np.float32)
            print("初始化数组sv_data成功！")

            # 计算lv和sv
            valid_mask = ~np.isnan(slop_data)
            lv_data[valid_mask], sv_data[valid_mask] = get_LS(slop_data[valid_mask])

            # 定义写入文件的参数
            rasterio_params = {
                'driver': 'GTiff',
                'height': slop_data.shape[0],
                'width': slop_data.shape[1],
                'count': 1,
                'dtype': rasterio.float32,
                'crs': 'EPSG:4326',  # src.crs
                'transform': src.transform,
                'nodata': -9999,
                'compress': 'lzw',
            }

            # 写入lv.tif
            with rasterio.open(output_lv_path, 'w', **rasterio_params) as dst:
                dst.write(lv_data, 1)

            # 写入sv.tif
            with rasterio.open(output_sv_path, 'w', **rasterio_params) as dst:
                dst.write(sv_data, 1)


    except Exception as e:
        print(f"处理文件 {input_file_path} 时发生错误: {e}")


def process_all_tifs(input_folder: str) -> None:
    """
    遍历指定文件夹中的所有tif文件，并对每个文件计算L和S值。

    参数:
    input_folder (str): 包含tif文件的文件夹路径。
    """
    for filename in os.listdir(input_folder):
        if filename.endswith('.tif'):
            print("正在处理的文件：", filename)
            input_file_path = os.path.join(input_folder, filename)
            output_lv_path = input_file_path.replace('slop', 'SR_L')
            output_sv_path = input_file_path.replace('slop', 'SR_S')

            # 检查文件是否存在
            if os.path.exists(output_lv_path) and os.path.exists(output_sv_path):
                print(f"文件 {output_lv_path} 和 {output_sv_path} 已存在，跳过保存操作。")
            else:
                # 确保输出路径存在，如果不存在则创建
                os.makedirs(os.path.dirname(output_lv_path), exist_ok=True)
                os.makedirs(os.path.dirname(output_sv_path), exist_ok=True)
                try:
                    calculate_lv_sv(input_file_path, output_lv_path, output_sv_path)
                    print(f"Processed {input_file_path} -> {output_lv_path} and {output_sv_path}")
                except Exception as e:
                    print(f"处理文件 {input_file_path} 时发生错误: {e}")


# 示例用法
input_folder = r'G:\GEP_data\slop'
process_all_tifs(input_folder)
"""
根据input_folder = r'H:\GEP_data\slop' 路径在H:\GEP_data中创建H:\GEP_data\SR_L和H:\GEP_data\SR_S来存放计算好的文件；
可以跳过已处理好的文件；


"""